from constants import *
from vm_codes import *


class CodeWriter():
    def __init__(self, file) -> None:
        self.file = open(file, 'w')
        self.label_index = 0
        self.current_function_name = ""
        self.writeInit()

    def __exit__(self, arg0, arg1, arg2):
        self.Close()

    def writeArithmetic(self, cmd):
        if cmd in ['add', 'sub', 'and', 'or']:
            cmd = two_operand_code(cmd)
        elif cmd in ['neg', 'not']:
            cmd = one_operand_code(cmd)
        elif cmd in ['eq', 'lt', 'gt']:
            cmd = jmp_code(cmd, get_new_label(self.label_index))
            self.label_index = self.label_index + 1
        self.file.write(cmd)

    def writePushPop(self, srcfile, cmd_type, segment, offset):
        if cmd_type == C_PUSH:
            self.file.write(push_code(srcfile, segment, offset))
        elif cmd_type == C_POP:
            self.file.write(pop_code(srcfile, segment, offset))

    def writeInit(self):
        self.file.write(init_code())

    def writeLabel(self, label):
        label = self.current_function_name + '$' + label
        self.file.write(label_code(label))

    def writeGoto(self, label):
        label = self.current_function_name + '$' + label
        self.file.write(goto_code(label))

    def writeIf(self, label):
        label = self.current_function_name + '$' + label
        self.file.write(if_goto_code(label))

    # push return_address
    # push LCL
    # push ARG
    # push THIS
    # push THAT
    # ARG = SP-n-5
    # LCL = SP
    # goto f
    # (return_address)

    def writeCall(self, functionName, numArgs):
        self.file.write(call_code(functionName, numArgs))

    # FRAME = LCL
    # RET = *(FRAME-5)
    # *ARG = pop()
    # SP = ARG+1
    # THAT = *(FRAME-1)
    # THIS = *(FRAME-2)
    # ARG = *(FRAME-3)
    # LCL = *(FRAME-4)
    # goto RET
    def writeReturn(self):
        self.file.write(return_code())

    # (f)
    # repeat k times: 初始化local段
    # PUSH 0
    def writeFunction(self, functionName, numLocals):
        self.current_function_name = functionName
        self.file.write(function_code(functionName, numLocals))

    def Close(self):
        self.file.close()
